#pragma once

#include <algorithm/math.h>
#include <initializer_list>
#include <assert.h>

namespace os {
namespace math {
template <size_t rowNbr, size_t colNbr>
class Matrix;
template <size_t size>
constexpr Matrix<size, size> getIdentityMatrix();
template <size_t size>
constexpr Matrix<size, size> MatrixI = getIdentityMatrix<size>();
} // namespace math
} // namespace os

namespace os {
namespace math {
template <size_t rowNbr, size_t colNbr>
class Matrix {
public:
    constexpr static size_t ROW = rowNbr;
    constexpr static size_t COL = colNbr;

    template <size_t friendRow, size_t friendCol>
    friend class Matrix;

private:
    float32_t _data[rowNbr * colNbr]{};

public:
    constexpr Matrix() = default;

    explicit constexpr Matrix(float32_t const initData[rowNbr * colNbr]) {
        for (size_t i = 0; i < rowNbr * colNbr; i++) {
            _data[i] = initData[i];
        }
    }

    constexpr Matrix(std::initializer_list<float32_t> initList) {
        assert(initList.size() == rowNbr * colNbr);
        uint16_t i = 0;
        for (auto& item : initList) {
            _data[i++] = item;
        }
    }

    constexpr Matrix(Matrix<rowNbr, colNbr> const& other) {
        for (size_t i = 0; i < rowNbr * colNbr; i++) {
            _data[i] = other._data[i];
        }
    }

    constexpr Matrix& operator=(float32_t const initData[rowNbr * colNbr]) {
        for (size_t i = 0; i < rowNbr * colNbr; i++) {
            _data[i] = initData[i];
        }
        return *this;
    }

    constexpr Matrix& operator=(Matrix<rowNbr, colNbr> const& other) noexcept {
        if (this != &other) {
            for (size_t i = 0; i < rowNbr * colNbr; i++) {
                _data[i] = other._data[i];
            }
        }
        return *this;
    }

    constexpr Matrix& operator=(std::initializer_list<float32_t> initList) noexcept {
        assert(initList.size() == rowNbr * colNbr);
        size_t i = 0;
        for (auto& item : initList) {
            _data[i++] = item;
        }
        return *this;
    }

    constexpr Matrix<rowNbr, colNbr> operator+(Matrix<rowNbr, colNbr> const& rhs) const {
        Matrix<rowNbr, colNbr> result;

        float32_t const* pInA = _data; /* 矩阵 A */
        float32_t const* pInB = rhs._data; /* 矩阵 B */
        float32_t* pOut = result._data; /* 输出矩阵 */

        uint32_t numSamples = rowNbr * colNbr; /* 矩阵元素总数 */
        uint32_t blkCnt = numSamples >> 2U; /* 循环计数器 */

        /* 使用递减计数器，速度更快 */
        while (blkCnt > 0U) {
            /* 使用循环展开技术，粒度为 4 */
            *pOut++ = *pInA++ + *pInB++;
            *pOut++ = *pInA++ + *pInB++;
            *pOut++ = *pInA++ + *pInB++;
            *pOut++ = *pInA++ + *pInB++;

            blkCnt--;
        }

        /* 循环展开善后 */
        blkCnt = numSamples % 0x4U;

        while (blkCnt > 0U) {
            *pOut++ = *pInA++ + *pInB++;

            blkCnt--;
        }

        return result;
    }

    constexpr Matrix<rowNbr, colNbr> operator-(Matrix<rowNbr, colNbr> const& rhs) const {
        Matrix<rowNbr, colNbr> result;

        float32_t const* pInA = _data; /* 矩阵 A */
        float32_t const* pInB = rhs._data; /* 矩阵 B */
        float32_t* pOut = result._data; /* 输出矩阵 */

        uint32_t numSamples = rowNbr * colNbr; /* 矩阵元素总数 */
        uint32_t blkCnt = numSamples >> 2U; /* 循环计数器 */

        /* 使用递减计数器，速度更快 */
        while (blkCnt > 0U) {
            /* 使用循环展开技术，粒度为 4 */
            *pOut++ = *pInA++ - *pInB++;
            *pOut++ = *pInA++ - *pInB++;
            *pOut++ = *pInA++ - *pInB++;
            *pOut++ = *pInA++ - *pInB++;

            blkCnt--;
        }

        /* 循环展开善后 */
        blkCnt = numSamples % 0x4U;

        while (blkCnt > 0U) {
            *pOut++ = *pInA++ - *pInB++;

            blkCnt--;
        }

        return result;
    }

    constexpr Matrix<rowNbr, colNbr> operator-() const {
        Matrix<rowNbr, colNbr> result;

        float32_t const* pInA = _data; /* 矩阵 A */
        float32_t* pOut = result._data; /* 输出矩阵 */

        uint32_t numSamples = rowNbr * colNbr; /* 矩阵元素总数 */
        uint32_t blkCnt = numSamples >> 2U; /* 循环计数器 */

        /* 使用递减计数器，速度更快 */
        while (blkCnt > 0U) {
            /* 使用循环展开技术，粒度为 4 */
            *pOut++ = -*pInA++;
            *pOut++ = -*pInA++;
            *pOut++ = -*pInA++;
            *pOut++ = -*pInA++;

            blkCnt--;
        }

        /* 循环展开善后 */
        blkCnt = numSamples % 0x4U;

        while (blkCnt > 0U) {
            *pOut++ = -*pInA++;

            blkCnt--;
        }

        return result;
    }

    template <size_t colResult>
    constexpr Matrix<rowNbr, colResult> operator*(Matrix<colNbr, colResult> const& rhs) const {
        Matrix<rowNbr, colResult> result;

        float32_t const* pIn1 = _data; /* 矩阵 A */
        float32_t const* pIn2 = rhs._data; /* 矩阵 B */
        float32_t const* pInA = _data; /* 矩阵 A */
        float32_t const* pInB = rhs._data; /* 矩阵 B */
        float32_t* pOut = result._data; /* 输出矩阵 */
        float32_t* px = nullptr; /* 临时指针 */
        float32_t sum = 0; /* 累加器 */
        uint32_t colResultCnt = 0, i = 0U, rowCnt = rowNbr, colCnt = 0; /* 循环计数器和指针 */

        /* 行循环 */
        do {
            px = pOut + i;
            colResultCnt = colResult;
            pIn2 = rhs._data;

            /* 列循环 */
            do {
                sum = 0.0f;
                pIn1 = pInA;
                colCnt = colNbr >> 2U;

                /* 核心 */
                while (colCnt > 0U) {
                    sum += *pIn1++ * *pIn2;
                    pIn2 += colResult;

                    sum += *pIn1++ * *pIn2;
                    pIn2 += colResult;

                    sum += *pIn1++ * *pIn2;
                    pIn2 += colResult;

                    sum += *pIn1++ * *pIn2;
                    pIn2 += colResult;

                    colCnt--;
                }

                colCnt = colNbr % 0x4U;

                while (colCnt > 0U) {
                    sum += *pIn1++ * *pIn2;
                    pIn2 += colResult;

                    colCnt--;
                }

                *px++ = sum;

                colResultCnt--;

                pIn2 = pInB + (colResult - colResultCnt);
            } while (colResultCnt > 0U);

            /* 更新指针到下一行 */
            i = i + colResult;
            pInA = pInA + colNbr;

            rowCnt--;
        } while (rowCnt > 0U);

        return result;
    }

    constexpr Matrix<rowNbr, colNbr> operator*(float32_t const rhs) const {
        Matrix<rowNbr, colNbr> result;

        float32_t const* pIn = _data; /* 输入矩阵 */
        float32_t* pOut = result._data; /* 输出矩阵 */
        uint32_t numSamples = rowNbr * colNbr; /* 元素总数 */
        uint32_t blkCnt = numSamples >> 2U; /* 循环计数器 */

        while (blkCnt > 0U) {
            *pOut++ = (*pIn++) * rhs;
            *pOut++ = (*pIn++) * rhs;
            *pOut++ = (*pIn++) * rhs;
            *pOut++ = (*pIn++) * rhs;

            blkCnt--;
        }

        blkCnt = numSamples % 0x4U;

        while (blkCnt > 0U) {
            *pOut++ = (*pIn++) * rhs;

            blkCnt--;
        }

        return result;
    }

    friend Matrix<rowNbr, colNbr> operator*(float32_t const lhs, Matrix<rowNbr, colNbr> const& mat) {
        return mat * lhs;
    }

    constexpr Matrix<rowNbr, colNbr> operator/(float32_t const rhs) const {
        Matrix<rowNbr, colNbr> result;

        float32_t const* pIn = _data; /* 输入矩阵 */
        float32_t* pOut = result._data; /* 输出矩阵 */
        uint32_t numSamples = rowNbr * colNbr; /* 元素总数 */
        uint32_t blkCnt = numSamples >> 2U; /* 循环计数器 */

        while (blkCnt > 0U) {
            *pOut++ = (*pIn++) / rhs;
            *pOut++ = (*pIn++) / rhs;
            *pOut++ = (*pIn++) / rhs;
            *pOut++ = (*pIn++) / rhs;

            blkCnt--;
        }

        blkCnt = numSamples % 0x4U;

        while (blkCnt > 0U) {
            *pOut++ = (*pIn++) / rhs;

            blkCnt--;
        }

        return result;
    }

    template <size_t row, size_t col>
    constexpr float32_t get() const {
        static_assert((row >= 0 and row < rowNbr) and (col >= 0 and col < colNbr), "Out of range");
        return _data[row * colNbr + col];
    }

    constexpr float32_t operator()(size_t row, size_t col) const {
        assert(row < rowNbr && col < colNbr);
        return _data[row * colNbr + col];
    }

    constexpr float32_t* operator[](size_t row) {
        assert(row < rowNbr);
        return _data + row * colNbr;
    }

    constexpr void swapRow(size_t col_from, size_t row1, size_t row2) {
        int32_t _w = 0;
        float32_t tmp = 0;
        float32_t* dataI = _data + row1 * colNbr + (col_from);
        float32_t* dataJ = _data + row2 * colNbr + (col_from);
        int32_t const nb = colNbr - col_from;

        for (_w = 0; _w < nb; _w++) {
            tmp = *dataI;
            *dataI++ = *dataJ;
            *dataJ++ = tmp;
        }
    }

    /**
     * @brief 缩放行
     * @param col_from 起始列
     * @param scale 乘数
     * @param row 指定行
     * @details row *= scale
     * @return void
     */
    constexpr void scaleRow(size_t col_from, float32_t scale, size_t row) {
        int32_t _w = 0;
        float32_t* data = _data + row * colNbr + (col_from);
        int32_t const nb = colNbr - col_from;

        for (_w = 0; _w < nb; _w++) {
            *data++ *= scale;
        }
    }

    /**
     * @brief 缩放列
     * @param row_from 起始行
     * @param scale 乘数
     * @param col 指定列
     * @details col *= scale
     * @return void
     */
    constexpr void scaleCol(size_t row_from, float32_t scale, size_t col) {
        int32_t _w = 0;
        float32_t* data = _data;
        int32_t const nb = colNbr - row_from;
        data += col + colNbr * row_from;
        for (_w = 0; _w < nb; _w++) {
            *data *= scale;
            data += colNbr;
        }
    }

    /**
     * @brief 初等行变换 用某个数乘以某一行减到另一行中去
     * @param col_from 起始列
     * @param row1 行 1
     * @param scale 乘数
     * @param row2 行 2
     * @details row1 -= scale * row2
     * @return void
     */
    constexpr void masRow(size_t col_from, size_t row1, float32_t scale, size_t row2) {
        int32_t _w = 0;
        float32_t* dataA = _data + row1 * colNbr + col_from;
        float32_t* dataB = _data + row2 * colNbr + col_from;
        int32_t const nb = colNbr - col_from;

        for (_w = 0; _w < nb; _w++) {
            *dataA++ -= scale * *dataB++;
        }
    }

    /**
     * @brief 约当-高斯法求逆矩阵
     * @return
     */
    constexpr Matrix<rowNbr, colNbr> inverse() const {
        static_assert(rowNbr == colNbr, "Matrix must be square to be inversed");
        Matrix<rowNbr, colNbr> src(*this);
        Matrix<rowNbr, colNbr> result(MatrixI<rowNbr>);

        float32_t* pIn = src._data;
        float32_t* pOut = result._data;

        float32_t* pTmp = pOut;

        float32_t pivot = 0.0f;
        float32_t newPivot = 0.0f;
        uint32_t selectedRow = 0;
        uint32_t pivotRow = 0;
        uint32_t i = 0;
        uint32_t rowNb = 0;
        uint32_t flag = 0U;
        uint32_t column = 0;

        for (column = 0U; column < rowNbr; column++) {
            pivotRow = column;

            pTmp = src[column] + column;
            pivot = *pTmp;
            selectedRow = column;

            for (rowNb = column + 1; rowNb < rowNbr; rowNb++) {
                pTmp = src[rowNb] + column;
                newPivot = *pTmp;
                if (fabsf(newPivot) > fabsf(pivot)) {
                    selectedRow = rowNb;
                    pivot = newPivot;
                }
            }

            if ((pivot != 0.0f) && (selectedRow != column)) {
                src.swapRow(column, pivotRow, selectedRow);
                result.swapRow(0, pivotRow, selectedRow);

                flag = 1U;
            }

            if ((flag != 1U) && (pivot == 0.0f)) {
                MathErrorHandler("Matrix is singular");
                return result;
            }

            pivot = 1.0f / pivot;

            src.scaleRow(column, pivot, pivotRow);
            result.scaleRow(0, pivot, pivotRow);

            rowNb = 0;
            for (; rowNb < pivotRow; rowNb++) {
                pTmp = src[rowNb] + column;
                pivot = *pTmp;

                src.masRow(column, rowNb, pivot, pivotRow);
                result.masRow(0, rowNb, pivot, pivotRow);
            }

            for (rowNb = pivotRow + 1; rowNb < rowNbr; rowNb++) {
                pTmp = src[rowNb] + column;
                pivot = *pTmp;

                src.masRow(column, rowNb, pivot, pivotRow);
                result.masRow(0, rowNb, pivot, pivotRow);
            }
        }

        if ((flag != 1U) && (pivot == 0.0f)) {
            pIn = src._data;
            for (i = 0; i < rowNbr * rowNbr; i++) {
                if (pIn[i] != 0.0f) break;
            }

            if (i == rowNbr * rowNbr) {
                MathErrorHandler("Matrix is singular");
                return result;
            }
        }

        return result;
    }

    constexpr Matrix operator/(Matrix<colNbr, colNbr> const& rhs) const { return (*this) * (rhs.inverse()); }

    constexpr Matrix<rowNbr, colNbr> operator^(int const& power) const {
        static_assert(rowNbr == colNbr, "Matrix must be square to be powered");
        if (power > 1) {
            Matrix<rowNbr, colNbr> result(*this);
            for (int i = 1; i < power; i++) {
                result = result * *this;
            }
            return result;
        }
        if (power == 1) {
            return *this;
        }
        if (power == 0) {
            return MatrixI<rowNbr>;
        }
        if (power == -1) {
            return this->inverse();
        }
        if (power < 0) {
            return (this->inverse()) ^ (-power);
        }
        return *this;
    }

    constexpr Matrix<colNbr, rowNbr> T() const {
        Matrix<colNbr, rowNbr> result;
        float32_t const* pIn = _data;
        float32_t* pOut = result._data;
        float32_t* px = nullptr;
        uint16_t nRows = rowNbr;
        uint16_t nCols = colNbr;
        uint32_t col = 0;
        uint32_t row = nRows;
        uint32_t i = 0U;

        do {
            px = pOut + i;

            col = nCols;

            while (col > 0U) {
                *px = *pIn++;

                px += nRows;

                col--;
            }
            i++;
            row--;
        } while (row > 0U);

        return result;
    }

    constexpr Matrix<colNbr, rowNbr> llt() {
        static_assert(colNbr == rowNbr, "Matrix must be square to be LLT");
        Matrix<colNbr, rowNbr> result;

        uint32_t i = 0;
        uint32_t j = 0;
        uint32_t k = 0;
        uint32_t n = colNbr;
        float32_t invSqrtVj = 0;
        float32_t* pA = _data;
        float32_t* pG = result._data;

        for (i = 0; i < n; i++) {
            for (j = i; j < n; j++) {
                pG[j * n + i] = pA[j * n + i];

                for (k = 0; k < i; k++) {
                    pG[j * n + i] = pG[j * n + i] - pG[i * n + k] * pG[j * n + k];
                }
            }

            if (pG[i * n + i] <= 0.0f) {
                MathErrorHandler("Matrix is not positive definite");
                return result;
            }
            invSqrtVj = 1.0f / sqrtf(pG[i * n + i]);

            result.scaleCol(i, invSqrtVj, i);
        }

        return result;
    }
};

template <size_t size>
constexpr Matrix<size, size> getIdentityMatrix() {
    float data[size * size]{};
    for (size_t i = 0; i < size; ++i) {
        data[i * size + i] = 1.0f;
    }
    Matrix<size, size> mat(data);
    return mat;
}

template <size_t size>
constexpr Matrix<size, size> getZeroMatrix() {
    float data[size * size]{};
    Matrix<size, size> mat(data);
    return mat;
}
} // namespace math
} // namespace os